from langchain_core.messages import HumanMessage
from langchain_core.tools import BaseTool

from model import LLMFactory
from state import MessagesState
from util import get_current_time


class Agent:
    def __init__(
        self,
        tools: list[BaseTool],
        prompt_template: str,
        model_provider: str = None,
        temperature: float = None,
    ):
        self.model = LLMFactory.get_llm(model_provider, temperature)
        if tools:
            self.model = self.model.bind_tools(tools)
        self.prompt_template = prompt_template


class AsyncAgent(Agent):
    async def __call__(self, state: MessagesState):
        messages = state["messages"]
        first_human_message = messages[0].content

        prompt = HumanMessage(
            self.prompt_template.format(
                first_user_message=first_human_message, current_time=get_current_time()
            )
        )
        messages = [prompt] + messages[1:]

        response = await self.model.ainvoke(messages)

        return {"messages": [response]}


class SyncAgent(Agent):
    def __call__(self, state: MessagesState):
        messages = state["messages"]
        first_human_message = messages[0]

        prompt = HumanMessage(self.prompt_template.format(first_user_message=first_human_message))
        messages = [prompt] + messages[1:]

        response = self.model.invoke(messages)

        return {"message": [response]}
